import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;

public class Server {
    public static final String IMAGE_PATH = "/root/yolov5-b/data/images/input_image.jpg";  // 固定保存的图片路径和文件名
    public static final String RESULT_IMAGE_PATH = "/root/yolov5-b/result/result.jpg";  // 检测结果图片的固定路径
    public static final String RESULT_JSON_PATH = "/root/yolov5-b/result/result.json";  // 检测结果JSON文件的固定路径

    public static final String YOLO_CD_COMMAND = "cd /root/yolov5-b";  // 进入yolov5-b文件夹
    public static final String YOLO_DETECT_COMMAND = "python detect.py  --source data/images/input_image.jpg";  // 运行YOLO检测的命令

    //private static boolean isEnvActivated = false;  // 标志虚拟环境是否已启动

    public static void main(String[] args) {
        int port = 8080; // 服务器端口号
        try (ServerSocket serverSocket = new ServerSocket(port)) {
            System.out.println("服务器启动，等待客户端连接...");

            while (true) {
                try (Socket clientSocket = serverSocket.accept();
                     DataInputStream dis = new DataInputStream(clientSocket.getInputStream());
                     DataOutputStream dos = new DataOutputStream(clientSocket.getOutputStream())) {

                    // 接收客户端发送的图片并保存到固定路径
                    long imageSize = dis.readLong();
                    FileOutputStream fosImage = new FileOutputStream(IMAGE_PATH);
                    byte[] buffer = new byte[4096];
                    int bytesRead;
                    while (imageSize > 0 && (bytesRead = dis.read(buffer, 0, (int)Math.min(buffer.length, imageSize))) != -1) {
                        fosImage.write(buffer, 0, bytesRead);
                        imageSize -= bytesRead;
                    }
                    fosImage.close();
                    System.out.println("图片已接收并保存到固定路径。");

                    // 启动YOLO虚拟环境并运行目标检测
                    runYOLODetection();

                    // 读取检测结果图片并发送给客户端
                    File resultImage = new File(RESULT_IMAGE_PATH);
                    FileInputStream fisResultImage = new FileInputStream(resultImage);
                    dos.writeLong(resultImage.length());
                    while ((bytesRead = fisResultImage.read(buffer)) != -1) {
                        dos.write(buffer, 0, bytesRead);
                    }
                    fisResultImage.close();

                    // 读取检测结果JSON文件并发送给客户端
                    File resultJson = new File(RESULT_JSON_PATH);
                    FileInputStream fisResultJson = new FileInputStream(resultJson);
                    dos.writeLong(resultJson.length());
                    while ((bytesRead = fisResultJson.read(buffer)) != -1) {
                        dos.write(buffer, 0, bytesRead);
                    }
                    fisResultJson.close();

                    System.out.println("检测结果已发送给客户端。");

                } catch (IOException e) {
                    e.printStackTrace();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 启动YOLO虚拟环境并运行目标检测
    private static void runYOLODetection() throws IOException, InterruptedException {

//        //先进入yolov5-b文件夹
//        System.out.println("进入yolov5文件夹中...");
//        ProcessBuilder pbActivateEnv = new ProcessBuilder("bash", "-c", YOLO_CD_COMMAND);
//        Process processActivateEnv = pbActivateEnv.start();
//        processActivateEnv.waitFor();
//        //isEnvActivated = true;  // 标记虚拟环境已经启动
//
//
//        // 运行YOLO检测
//        System.out.println("运行YOLO目标检测...");
//        ProcessBuilder pbRunDetect = new ProcessBuilder("bash", "-c", YOLO_DETECT_COMMAND);
//        Process processRunDetect = pbRunDetect.start();
//        processRunDetect.waitFor();
//
//        System.out.println("目标检测完成。");
        System.out.println("进入yolov5文件夹并运行YOLO目标检测...");
        Runtime runtime = Runtime.getRuntime();
        Process process = runtime.exec(new String[]{"bash", "-c", YOLO_CD_COMMAND + " && " + YOLO_DETECT_COMMAND});

        // 输出命令执行过程中的日志
        BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
        String line;
        while ((line = reader.readLine()) != null) {
            System.out.println(line);
        }

        // 输出错误日志（如果有）
        BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream()));
        while ((line = errorReader.readLine()) != null) {
            System.err.println(line);
        }

        // 等待进程完成
        int exitCode = process.waitFor();
        if (exitCode != 0) {
            System.err.println("命令执行失败，退出代码：" + exitCode);
        }

        System.out.println("目标检测完成。");
    }
}
